import tensorflow as tf

# 可以通过 tf.config.experimental.set_memory_growth 将 GPU 的显存使用策略设置为 “仅在需要时申请显存空间”。以下代码将所有 GPU 设置为仅在需要时申请显存空间：

gpus = tf.config.list_physical_devices(device_type='GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(device=gpu, enable=True)

# 以下代码通过 tf.config.set_logical_device_configuration 选项并传入 tf.config.LogicalDeviceConfiguration 实例，设置 TensorFlow 固定消耗 GPU:0 的 1GB 显存（其实可以理解为建立了一个显存大小为 1GB 的 “虚拟 GPU”）：

gpus = tf.config.list_physical_devices(device_type='GPU')
tf.config.set_logical_device_configuration(
    gpus[0],
    [tf.config.LogicalDeviceConfiguration(memory_limit=1024)])
